﻿#if defined(_WIN32)

#include	"Engine.Network.hpp"
#include	"Engine.Memory.hpp"
#include	"Engine.Thread.hpp"

#include	<WinSock2.h>

#include	<atomic>
#include	<list>
#include	<map>
#include	<thread>

namespace Engine { namespace Network {	

	struct UserOverLapped : public OVERLAPPED {
		uint64_t	nConnId;
		int			nOpt;
		char *		pData;
	};

	typedef Memory::ObjectPool<UserOverLapped> OverLappedPool;

	struct Session {
		SOCKET			iSocket;		//! 连接Socket
		SOCKADDR_IN		iAddr;			//! Ip地址
		void *			pUserData;		//! 服务器为完成端句柄，客户端为UserOverLapped
	};

	class SocketImpl : public Socket {
	public:
		SocketImpl();
		virtual ~SocketImpl();

		virtual int		Connect(const char * pStrIp, unsigned short uPort, IListener * pListener);
		virtual int		Send(const char * pData, size_t uSize);
		virtual void	Close(int nReason = ECloseReason::ClientClose);

	private:
		void			_IOThread();

	private:
		std::atomic<bool>	_bValid;
		Session				_iSession;
		IListener *			_pListener;
		HANDLE				_hCompletePort;
		Thread::Thread		_iIOWorker;
		OverLappedPool		_iOverlappedPool;
	};

	SocketImpl::SocketImpl()
		: _bValid(false)
		, _iSession()
		, _pListener(nullptr)
		, _iIOWorker()
		, _iOverlappedPool(1024, 32) {
		WORD		wVersion = MAKEWORD(2, 2);
		WSADATA		wsaData;

		if (::WSAStartup(wVersion, &wsaData))
			throw std::logic_error("::WSAStartup Failed!");
	}

	SocketImpl::~SocketImpl() {
		Close();
		::WSACleanup();
	}

	int SocketImpl::Connect(const char * pStrIp, unsigned short uPort, IListener * pListener) {
		UserOverLapped * pData = nullptr;
		DWORD dwFlag = 0, dwReceived = 0;
		WSABUF wsaBuf;

		if (!pStrIp || strlen(pStrIp) <= 1 || uPort <= 0 || !pListener) return ErrParam;
		if (_bValid) return ErrAlreadyUsed;
				
		pData = _iOverlappedPool.Alloc();
		if (!pData)
			return ErrAlloc;

		pData->nConnId	= 1;
		pData->nOpt		= FD_READ;
		pData->pData		= new (std::nothrow) char[ReceiveBuffer];
		
		if (!pData->pData) {
			_iOverlappedPool.Free(pData);
			return ErrAlloc;
		}			

		ZeroMemory(pData->pData, ReceiveBuffer);
		_iSession.pUserData = pData;
		
		_iSession.iSocket = ::WSASocket(AF_INET, SOCK_STREAM, 0, 0, 0, WSA_FLAG_OVERLAPPED);
		if (_iSession.iSocket == INVALID_SOCKET) {
			delete[] pData->pData;
			_iOverlappedPool.Free(pData);
			return ::WSAGetLastError();
		}

		::ZeroMemory(&_iSession.iAddr, sizeof(_iSession.iAddr));

		_iSession.iAddr.sin_family = AF_INET;
		_iSession.iAddr.sin_port = htons(uPort);
		_iSession.iAddr.sin_addr.S_un.S_addr = inet_addr(pStrIp);

		if (::connect(_iSession.iSocket, (SOCKADDR *)&_iSession.iAddr, sizeof(SOCKADDR)) != 0) {
			delete[] pData->pData;
			_iOverlappedPool.Free(pData);
			return ::WSAGetLastError();
		}

		_hCompletePort = ::CreateIoCompletionPort(INVALID_HANDLE_VALUE, 0, 0, 0);
		if (!::CreateIoCompletionPort((HANDLE)_iSession.iSocket, _hCompletePort, 0, 0)) {
			_bValid = false;
			delete[] pData->pData;
			_iOverlappedPool.Free(pData);
			::closesocket(_iSession.iSocket);
			return ::GetLastError();
		}

		_bValid = true;

		if (!_iIOWorker.Start(&SocketImpl::_IOThread, this)) {
			_bValid = false;
			delete[] pData->pData;
			_iOverlappedPool.Free(pData);
			::closesocket(_iSession.iSocket);
			::CloseHandle(_hCompletePort);
			return ErrCreateThread;
		}

		_pListener = pListener;
		_pListener->pSender = this;
		_pListener->OnConnect(_iSession.iAddr.sin_addr.S_un.S_addr, uPort);		
		
		wsaBuf.buf = pData->pData;
		wsaBuf.len = ReceiveBuffer;
		
		if (SOCKET_ERROR == ::WSARecv(_iSession.iSocket, &wsaBuf, 1, &dwReceived, &dwFlag, pData, 0)) {
			if (::WSAGetLastError() != ERROR_IO_PENDING) {
				Close(ECloseReason::DataError);
				return ::WSAGetLastError();
			}
		}

		return Success;
	}

	int SocketImpl::Send(const char * pData, size_t uSize) {
		if (!_bValid) return ErrInvalidConnId;

		UserOverLapped * pOver = _iOverlappedPool.Alloc();
		if (!pOver)
			return ErrAlloc;

		char * pBuffer = new (std::nothrow) char[uSize];
		if (!pBuffer)
			return ErrAlloc;

		::memcpy(pBuffer, pData, uSize);
		
		pOver->nConnId	= 1;
		pOver->nOpt		= FD_WRITE;
		pOver->pData 	= pBuffer;

		WSABUF wsaBuf;
		wsaBuf.buf = pBuffer;
		wsaBuf.len = uSize;

		DWORD dwSend;
		if (SOCKET_ERROR == ::WSASend(_iSession.iSocket, &wsaBuf, 1, &dwSend, 0, pOver, 0)) {
			if (::WSAGetLastError() != ERROR_IO_PENDING) {
				_iOverlappedPool.Free(pOver);
				delete[] pBuffer;
				return ::WSAGetLastError();
			}
		}

		return Success;
	}

	void SocketImpl::Close(int nReason) {
		if (!_bValid) return;

		_bValid = false;

		::closesocket(_iSession.iSocket);
		::CloseHandle(_hCompletePort);

		UserOverLapped * pData = (UserOverLapped *) _iSession.pUserData;
		if (pData) {
			if (pData->pData) delete[] pData->pData;
			_iOverlappedPool.Free(pData);
		}

		_iOverlappedPool.FreeAll();
		_pListener->OnClose(nReason);
	}

	void SocketImpl::_IOThread() {
		DWORD				dwNumOfBytes;
		DWORD				dwNonUsed;
		INT					nRet;
		UserOverLapped *	pOverlapped;

		while (_bValid) {
			nRet = ::GetQueuedCompletionStatus(
				_hCompletePort,
				&dwNumOfBytes,
				(PULONG_PTR)&dwNonUsed,
				(LPOVERLAPPED *)&pOverlapped,
				WSA_INFINITE);

			if (!pOverlapped) continue;
			
			if (!nRet) {
				Close(ECloseReason::SystemClose);
				break;
			}

			if (dwNumOfBytes == 0 && (pOverlapped->nOpt == FD_READ || pOverlapped->nOpt == FD_WRITE)) {
				Close(ECloseReason::SystemClose);
				break;
			}

			switch (pOverlapped->nOpt) {
			case FD_READ:
				{
					_pListener->OnReceive(pOverlapped->pData, dwNumOfBytes);
					ZeroMemory(pOverlapped->pData, ReceiveBuffer);

					DWORD dwFlag = 0, dwReceived = 0;
					WSABUF wsaBuf;
					wsaBuf.buf = pOverlapped->pData;
					wsaBuf.len = ReceiveBuffer;

					if (SOCKET_ERROR == ::WSARecv(_iSession.iSocket, &wsaBuf, 1, &dwReceived, &dwFlag, pOverlapped, 0)) {
						if (::WSAGetLastError() != ERROR_IO_PENDING) {
							Close(ECloseReason::DataError);
							return;
						}
					}
				}
				break;
			case FD_WRITE:
				{
					delete[] (char *)pOverlapped->pData;
					_iOverlappedPool.Free(pOverlapped);
				}
				break;
			default:
				break;
			}
		}
	}

	class ServerSocketImpl : public ServerSocket {
	public:
		ServerSocketImpl();
		virtual ~ServerSocketImpl();

		virtual int		Listen(const char * pStrIp, unsigned short uPort, IListener * pListener);
		virtual int		Send(uint64_t nConnId, const char * pData, size_t uSize);
		virtual int		Broadcast(const char * pData, size_t uSize);
		virtual void	Shutdown();
		virtual void	Close(uint64_t nConnId, int nReason);

	private:
		void			_AcceptThread();			//! 处理连接请求的工作线程
		void			_IOThread();				//! 处理通信请求的工作线程

	private:
		std::atomic<bool>				_bValid;
		uint64_t						_nNextConnId;
		IListener *						_pListener;
		Session							_iServerSession;
		std::map<uint64_t, Session>		_mpClient;
		std::thread *					_pAcceptWorker;
		std::list<std::thread *>		_lIOWorker;
		OverLappedPool					_iOverLappedPool;
	};

	ServerSocketImpl::ServerSocketImpl()
		: _bValid(false)
		, _nNextConnId(0)
		, _pListener(nullptr)
		, _iServerSession()
		, _mpClient()
		, _pAcceptWorker(nullptr)
		, _lIOWorker()
		, _iOverLappedPool(10240, 64) {
		WORD		wVersion = MAKEWORD(2, 2);
		WSADATA		wsaData;

		if (::WSAStartup(wVersion, &wsaData))
			throw std::logic_error("::WSAStartup Failed!");
	}

	ServerSocketImpl::~ServerSocketImpl() {
		Shutdown();
		::WSACleanup();
	}

	int ServerSocketImpl::Listen(const char * pStrIp, unsigned short uPort, IListener * pListener) {
		if (!pStrIp || strlen(pStrIp) <= 1 || uPort <= 0 || !pListener) return ErrParam;
		if (_bValid) return ErrAlreadyUsed;

		SOCKET iServerSocket = ::WSASocket(AF_INET, SOCK_STREAM, 0, 0, 0, WSA_FLAG_OVERLAPPED);
		if (iServerSocket == INVALID_SOCKET)
			return ::WSAGetLastError();

		SOCKADDR_IN	iServerAddr;
		::ZeroMemory(&iServerAddr, sizeof(iServerAddr));

		iServerAddr.sin_family = AF_INET;
		iServerAddr.sin_port = htons(uPort);
		iServerAddr.sin_addr.S_un.S_addr = inet_addr(pStrIp);

		if (::bind(iServerSocket, (SOCKADDR *)&iServerAddr, sizeof(SOCKADDR)) != 0)
			return ::WSAGetLastError();

		if (::listen(iServerSocket, SOMAXCONN) != 0)
			return ::WSAGetLastError();

		SYSTEM_INFO iInfo;
		::ZeroMemory(&iInfo, sizeof(iInfo));
		::GetSystemInfo(&iInfo);

		_bValid = true;

		HANDLE hCompletion = ::CreateIoCompletionPort(INVALID_HANDLE_VALUE, 0, 0, 0);
		for (size_t nLoop = 0; nLoop < iInfo.dwNumberOfProcessors * 2; ++nLoop) {
			std::thread * pIOWorker = new (std::nothrow) std::thread(&ServerSocketImpl::_IOThread, this);
			if (!pIOWorker) {
				Shutdown();
				return ErrCreateThread;
			}
			_lIOWorker.push_back(pIOWorker);
		}

		_pListener = pListener;
		_pListener->pSender = this;

		_iServerSession.iSocket = iServerSocket;
		_iServerSession.pUserData = hCompletion;
		_iServerSession.iAddr = iServerAddr;

		_pAcceptWorker = new (std::nothrow) std::thread(&ServerSocketImpl::_AcceptThread, this);
		if (!_pAcceptWorker) {
			Shutdown();
			return ErrCreateThread;
		}

		return Success;
	}

	int ServerSocketImpl::Send(uint64_t nConnId, const char * pData, size_t uSize) {
		if (!_bValid || nConnId > _nNextConnId || _mpClient.find(nConnId) == _mpClient.end() || !pData || uSize <= 0)
			return ErrParam;

		Session iSession = _mpClient[nConnId];
		UserOverLapped * pOver = _iOverLappedPool.Alloc();
		if (!pOver)
			return ErrAlloc;

		char * pBuffer = new (std::nothrow) char[uSize];
		if (!pBuffer)
			return ErrAlloc;

		::memcpy(pBuffer, pData, uSize);
		
		pOver->nConnId	= nConnId;
		pOver->nOpt		= FD_WRITE;
		pOver->pData 	= pBuffer;

		WSABUF wsaBuf;
		wsaBuf.buf = pBuffer;
		wsaBuf.len = uSize;

		DWORD dwSend;
		if (SOCKET_ERROR == ::WSASend(iSession.iSocket, &wsaBuf, 1, &dwSend, 0, pOver, 0)) {
			if (::WSAGetLastError() != ERROR_IO_PENDING) {
				_iOverLappedPool.Free(pOver);
				delete[] pBuffer;
				return ::WSAGetLastError();
			}
		}

		return Success;
	}

	int ServerSocketImpl::Broadcast(const char * pData, size_t uSize) {
		if (!_bValid)
			return ErrInvalidConnId;

		char * pBuffer = new (std::nothrow) char[uSize];
		if (!pBuffer) return ErrAlloc;
		else ::memcpy(pBuffer, pData, uSize);

		for (auto itor = _mpClient.begin(); itor != _mpClient.end(); ++itor) {
			Session iSession		= itor->second;
			UserOverLapped * pOver	= _iOverLappedPool.Alloc();
			if (!pOver) return ErrAlloc;

			pOver->nConnId	= itor->first;
			pOver->nOpt		= FD_WRITE;

			WSABUF wsaBuf;
			wsaBuf.buf = pBuffer;
			wsaBuf.len = uSize;

			DWORD dwSend;
			if (SOCKET_ERROR == ::WSASend(iSession.iSocket, &wsaBuf, 1, &dwSend, 0, pOver, 0)) {
				if (::WSAGetLastError() != ERROR_IO_PENDING) {
					_iOverLappedPool.Free(pOver);
					delete[] pBuffer;
					return ::WSAGetLastError();
				}
			}
		}

		return Success;
	}

	void ServerSocketImpl::Close(uint64_t nConnId, int nReason) {
		if (nConnId > _nNextConnId)
			return;

		auto itor = _mpClient.find(nConnId);
		if (itor == _mpClient.end())
			return;

		_pListener->OnClose(nConnId, nReason);

		::closesocket(itor->second.iSocket);

		UserOverLapped * pData = (UserOverLapped *)itor->second.pUserData;
		delete[] pData->pData;
		_iOverLappedPool.Free(pData);

		_mpClient.erase(itor);
	}

	void ServerSocketImpl::Shutdown() {
		std::list<uint64_t> vConnId;
		
		for (auto itor = _mpClient.begin(); itor != _mpClient.end(); ++itor)
			vConnId.push_back(itor->first);

		for (auto itor = vConnId.begin(); itor != vConnId.end(); ++itor)
			Close(*itor, -1);
		
		if (!_bValid) return;
		_bValid = false;

		::closesocket(_iServerSession.iSocket);
		::CloseHandle(_iServerSession.pUserData);

		for (auto itor = _lIOWorker.begin(); itor != _lIOWorker.end(); ++itor) {
			std::thread * pWorker = *itor;
			if (pWorker->joinable()) pWorker->join();
			delete pWorker;
		}

		if (_pAcceptWorker) {
			if (_pAcceptWorker->joinable()) _pAcceptWorker->join();
			delete _pAcceptWorker;
			_pAcceptWorker = nullptr;
		}	

		_iOverLappedPool.FreeAll();	
	}

	void ServerSocketImpl::_AcceptThread() {
		while (_bValid) {
			SOCKADDR_IN iClientAddr;
			int	nSocketSize = sizeof(iClientAddr);

			SOCKET iClient = ::accept(_iServerSession.iSocket, (SOCKADDR *)&iClientAddr, &nSocketSize);
			if (iClient == INVALID_SOCKET) continue;
			
			uint64_t nConnId = ++_nNextConnId;
			UserOverLapped * pData = _iOverLappedPool.Alloc();

			if (!pData)	{
				::closesocket(iClient);
				_pListener->OnAccept(nConnId, iClientAddr.sin_addr.S_un.S_addr, iClientAddr.sin_port, ErrAlloc);
				continue;
			}
			
			pData->nConnId	= nConnId;
			pData->nOpt		= FD_READ;
			pData->pData		= new (std::nothrow) char[ReceiveBuffer];

			if (!pData->pData) {
				::closesocket(iClient);
				_iOverLappedPool.Free(pData);
				_pListener->OnAccept(nConnId, iClientAddr.sin_addr.S_un.S_addr, iClientAddr.sin_port, ErrAlloc);
				continue;
			}

			ZeroMemory(pData->pData, ReceiveBuffer);

			Session iClientSession;
			iClientSession.iSocket = iClient;
			iClientSession.iAddr = iClientAddr;
			iClientSession.pUserData = pData;

			WSABUF wsaBuf;
			wsaBuf.buf = pData->pData;
			wsaBuf.len = ReceiveBuffer;

			if (::CreateIoCompletionPort((HANDLE)iClient, _iServerSession.pUserData, 0, 0)) {
				_mpClient.insert(std::make_pair(pData->nConnId, iClientSession));

				DWORD dwFlag = 0, dwReceived = 0;
				if (SOCKET_ERROR == ::WSARecv(iClient, &wsaBuf, 1, &dwReceived, &dwFlag, pData, 0)) {
					if (::WSAGetLastError() != ERROR_IO_PENDING) {
						_pListener->OnAccept(nConnId, iClientAddr.sin_addr.S_un.S_addr, iClientAddr.sin_port, ::WSAGetLastError());
						Close(pData->nConnId, ECloseReason::DataError);
						continue;
					}
				}
				
				_pListener->OnAccept(nConnId, iClientAddr.sin_addr.S_un.S_addr, iClientAddr.sin_port, Success);
			} else {
				delete[] pData->pData;
				_iOverLappedPool.Free(pData);
				::closesocket(iClient);
				_pListener->OnAccept(nConnId, iClientAddr.sin_addr.S_un.S_addr, iClientAddr.sin_port, ::GetLastError());
			}
		}
	}

	void ServerSocketImpl::_IOThread() {
		DWORD				dwNumOfBytes;
		DWORD				dwNonUsed;
		INT					nRet;
		UserOverLapped *	pOverlapped;

		while (_bValid) {
			nRet = ::GetQueuedCompletionStatus(
				_iServerSession.pUserData,
				&dwNumOfBytes,
				(PULONG_PTR)&dwNonUsed,
				(LPOVERLAPPED *)&pOverlapped,
				WSA_INFINITE);
			
			if (!pOverlapped) continue;

			auto itor = _mpClient.find(pOverlapped->nConnId);
			if (itor == _mpClient.end())
				continue;

			if (!nRet) {
				Close(pOverlapped->nConnId, ECloseReason::DataError);
				continue;
			}

			if (dwNumOfBytes == 0 && (pOverlapped->nOpt == FD_READ || pOverlapped->nOpt == FD_WRITE)) {
				Close(pOverlapped->nConnId, ECloseReason::ClientClose);
				continue;
			}			

			Session iSession = itor->second;
			
			switch (pOverlapped->nOpt) {
			case FD_READ:
				{
					_pListener->OnReceive(pOverlapped->nConnId, pOverlapped->pData, dwNumOfBytes);
					ZeroMemory(pOverlapped->pData, ReceiveBuffer);

					WSABUF wsaBuf;
					wsaBuf.buf = pOverlapped->pData;
					wsaBuf.len = ReceiveBuffer;

					DWORD dwFlag = 0, dwReceived = 0;
					if (SOCKET_ERROR == ::WSARecv(iSession.iSocket, &wsaBuf, 1, &dwReceived, &dwFlag, pOverlapped, 0)) {
						if (::WSAGetLastError() != ERROR_IO_PENDING) {
							Close(pOverlapped->nConnId, ECloseReason::DataError);
						}
					}
				}
				break;
			case FD_WRITE:
				{
					delete[] (char *)pOverlapped->pData;
					_iOverLappedPool.Free(pOverlapped);
				}
				break;
			default:
				break;
			}
		}
	}

	Socket * Socket::Create() {
		return new (std::nothrow) SocketImpl();
	}

	void Socket::Close(Socket * pClient, int nReason) {
		if (!pClient) return;
		SocketImpl * pConn = (SocketImpl *)pClient;
		delete pConn;
	}

	ServerSocket * ServerSocket::Create() {
		return new (std::nothrow) ServerSocketImpl();
	}

	void ServerSocket::Shutdown(ServerSocket * pServer) {
		if (!pServer) return;
		ServerSocketImpl * pImpl = (ServerSocketImpl *) pServer;
		delete pImpl;
	}

}}

#endif
